package 笔试;

import java.util.*;

public class 与目标颜色最近的距离 {

    public List<Integer> shortestDistanceColorNB(int[] colors, int[][] queries) {
        HashMap<Integer, TreeSet<Integer>> map = new HashMap<>();
        for(int i = 0; i < colors.length; i++) {
            map.putIfAbsent(colors[i], new TreeSet<Integer>());
            map.get(colors[i]).add(i);
        }

        List<Integer> res = new ArrayList<>();
        for(int i = 0; i < queries.length; i++) {
            if(!map.containsKey(queries[i][1])) {
                res.add(-1);
                continue;
            }
            TreeSet<Integer> set = map.get(queries[i][1]);
            int dis = Integer.MAX_VALUE;
            // 和 queries[i][0] 位置最近的 queries[i][1] 的位置，从map里取set
            // 找右边, ceiling: 返回此集合中大于给定元素的最小元素
            if(set.ceiling(queries[i][0]) != null) {
                dis = Math.abs(queries[i][0] - (set.ceiling(queries[i][0])));
            }
            // 找左边, floor: 返回集合中小于给定元素的最大元素
            if(set.floor(queries[i][0]) != null) {
                dis = Math.min(dis, Math.abs(queries[i][0] - (set.floor(queries[i][0]))));
            }
            if(dis == Integer.MAX_VALUE) {
                res.add(-1);
            } else {
                res.add(dis);
            }
        }
        return res;
    }

    public List<Integer> shortestDistanceColor(int[] colors, int[][] queries) {
        int n=colors.length;
        int []first=new int[n];
        int []second=new int[n];
        int []third=new int[n];

        first[0]=colors[0]==1?0:n+1;
        second[0]=colors[0]==2?0:n+1;
        third[0]=colors[0]==3?0:n+1;
        for(int i=1;i<n;i++){
            first[i]=colors[i]==1?0:first[i-1]+1;
            second[i]=colors[i]==2?0:second[i-1]+1;
            third[i]=colors[i]==3?0:third[i-1]+1;
        }
        for(int i=n-2;i>=0;i--){
            first[i]=Math.min(first[i],first[i+1]+1);
            second[i]=Math.min(second[i],second[i+1]+1);
            third[i]=Math.min(third[i],third[i+1]+1);
        }

        List<Integer> ans=new LinkedList<>();
        for(int i=0;i<queries.length;i++){
            int t=-1;
            switch(queries[i][1]){
                case(1):t=first[queries[i][0]];break;
                case(2):t=second[queries[i][0]];break;
                case(3):t=third[queries[i][0]];break;
            }
            if(t>n) t=-1;
            ans.add(t);
        }

        return ans;
    }

    public List<Integer> shortestDistanceColor1(int[] colors, int[][] queries) {
        HashMap<Integer, TreeSet<Integer>> map = new HashMap<>();
        for(int i = 0; i < colors.length; i++)
        {
            map.putIfAbsent(colors[i], new TreeSet<Integer>());
            map.get(colors[i]).add(i);
        }
        List<Integer> res = new ArrayList<>();
        for(int i = 0; i < queries.length; i++)
        {
            if(!map.containsKey(queries[i][1]))
            {
                res.add(-1);
                continue;
            }

            TreeSet<Integer> set = map.get(queries[i][1]);
            int dis = Integer.MAX_VALUE;
            if(set.ceiling(queries[i][0]) != null)
                dis = Math.abs(set.ceiling(queries[i][0]) - queries[i][0]);
            if(set.floor(queries[i][0]) != null)
                dis = Math.min(dis, Math.abs(set.floor(queries[i][0]) - queries[i][0]));
            if(dis != Integer.MAX_VALUE)
                res.add(dis);
            else
                res.add(-1);
        }
        return res;
    }
}
